# 引入 PyTorch 相关模块
import torch
from torch.utils.data import DataLoader  # 用于批量读取数据

# 引入我们定义的三种模型（这里只用其中之一）
from models.cnn_model import TextCNN
from models.rnn_model import TextRNN
from models.attention_rnn_model import AttentionRNN

# 导入工具函数：用于加载数据集和数据处理类
from utils import load_data, TextDataset

# 用于绘图的库
import matplotlib.pyplot as plt

# ==================== 1. 训练参数设置 ====================
embed_dim = 100         # 词向量的维度（一个词用多少个数字表示）
hidden_dim = 128        # 隐藏层维度（仅对RNN类模型有用）
num_classes = 2         # 分类的类别数量（这里是二分类：0或1）
batch_size = 16         # 每批次喂入多少条数据
epochs = 10             # 训练多少轮
learning_rate = 1e-3    # 学习率（模型学习的“步子”大小）

# 选择设备：有GPU就用GPU，加速训练；没有就用CPU
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')


# ==================== 2. 加载数据 ====================
# 加载处理好的训练数据、验证数据和词汇表
train_data, val_data, vocab = load_data()

# 将数据封装成 PyTorch 的 Dataset，并使用 DataLoader 进行批处理
train_loader = DataLoader(TextDataset(train_data, vocab), batch_size=batch_size, shuffle=True)
val_loader = DataLoader(TextDataset(val_data, vocab), batch_size=batch_size)


# ==================== 3. 定义模型、优化器和损失函数 ====================
# 这里选择使用 AttentionRNN 模型（你也可以改成 TextCNN 或 TextRNN）
# len(vocab) 是词表大小，也就是 Embedding 层的输入维度
# model = AttentionRNN(len(vocab), embed_dim, hidden_dim, num_classes).to(device)


model = TextRNN(len(vocab), embed_dim, hidden_dim, num_classes).to(device)

# 使用 Adam 优化器（是一种常用的自动调节学习率的优化算法）
optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)

# 损失函数：交叉熵，用于分类问题
criterion = torch.nn.CrossEntropyLoss()


# 用来记录每轮训练和验证的准确率、损失
train_acc, val_acc = [], []
train_loss, val_loss = [], []


# ==================== 4. 模型训练 ====================
for epoch in range(epochs):
    # ---- 开始训练模式（启用 Dropout 等） ----
    model.train()

    correct = total = running_loss = 0  # 初始化本轮统计值

    # 遍历训练数据集
    for x, y in train_loader:
        # 将数据转移到 GPU 或 CPU
        x, y = x.to(device), y.to(device)

        optimizer.zero_grad()         # 清空梯度
        output = model(x)             # 前向传播
        loss = criterion(output, y)   # 计算损失
        loss.backward()               # 反向传播
        optimizer.step()              # 更新模型参数

        running_loss += loss.item()   # 累加损失值
        _, predicted = output.max(1)  # 取出预测结果（最大概率的类别）
        total += y.size(0)            # 累加总样本数
        correct += predicted.eq(y).sum().item()  # 累加预测正确的数量

    # 一轮训练结束，记录准确率和平均损失
    train_acc.append(correct / total)
    train_loss.append(running_loss / len(train_loader))


    # ---- 验证模式（不更新参数） ----
    model.eval()
    correct = total = val_running_loss = 0

    # 不需要计算梯度，提高效率
    with torch.no_grad():
        for x, y in val_loader:
            x, y = x.to(device), y.to(device)
            output = model(x)
            loss = criterion(output, y)

            val_running_loss += loss.item()
            _, predicted = output.max(1)
            total += y.size(0)
            correct += predicted.eq(y).sum().item()

    # 一轮验证结束，记录准确率和平均损失
    val_acc.append(correct / total)
    val_loss.append(val_running_loss / len(val_loader))

    # 输出每一轮训练和验证的结果
    print(f'Epoch [{epoch+1}/{epochs}], Train Acc: {train_acc[-1]:.4f}, Val Acc: {val_acc[-1]:.4f}')


# ==================== 5. 绘制准确率和损失曲线 ====================
plt.figure(figsize=(10,4))

# 左图：绘制准确率
plt.subplot(1,2,1)
plt.plot(train_acc, label='Train Acc')   # 训练准确率
plt.plot(val_acc, label='Val Acc')       # 验证准确率
plt.legend()
plt.title('Accuracy')

# 右图：绘制损失
plt.subplot(1,2,2)
plt.plot(train_loss, label='Train Loss') # 训练损失
plt.plot(val_loss, label='Val Loss')     # 验证损失
plt.legend()
plt.title('Loss')

# 保存图像为文件
plt.savefig('curve.png')
plt.show()


# ==================== 6. 保存训练好的模型参数 ====================
# 只保存模型的权重参数（不含结构）
torch.save(model.state_dict(), 'model.pth')
print("✅ 模型已保存为 model.pth")
